import json
import asyncio
import time

from fastapi import Body
from sse_starlette.sse import EventSourceResponse
from configs import LLM_MODELS, TEMPERATURE, HISTORY_LEN, Agent_MODEL

from langchain.chains import LLMChain
from langchain.memory import ConversationBufferWindowMemory
from langchain.agents import LLMSingleActionAgent, AgentExecutor
from typing import AsyncIterable, Optional, List

from server.utils import wrap_done, get_ChatOpenAI, get_prompt_template
from server.knowledge_base.kb_service.base import get_kb_details
from server.agent.custom_agent.ChatGLM3Agent import initialize_glm3_agent
from server.agent.tools_select import tools, tool_names
from server.agent.callbacks import CustomAsyncIteratorCallbackHandler, Status
from server.chat.utils import History, log_content
from server.agent import model_container
from server.agent.custom_template import CustomOutputParser, CustomPromptTemplate
from server.utils import create_chat_response, ChatMessageContent, ChatResponseCompletions, jsonify
from server.redis.redis_utils import ChatHistory, RedisClient
from startup import REDIS


async def agent_chat(uid: str = Body(..., description="用户ID"),
                     query: str = Body(..., description="用户输入", examples=["恼羞成怒"]),
                     unique_id: str = Body(..., description="每轮对话的唯一标识符"),
                     history: List[History] = Body([],
                                                   description="历史对话",
                                                   examples=[[
                                                       {"role": "user", "content": "请使用知识库工具查询今天北京天气"},
                                                       {"role": "assistant",
                                                        "content": "使用天气查询工具查询到今天北京多云，10-14摄氏度，东北风2级，易感冒"}]]
                                                   ),
                     stream: bool = Body(False, description="流式输出"),
                     model_name: str = Body(LLM_MODELS[0], description="LLM 模型名称。"),
                     temperature: float = Body(TEMPERATURE, description="LLM 采样温度", ge=0.0, le=1.0),
                     max_tokens: Optional[int] = Body(None, description="限制LLM生成Token数量，默认None代表模型最大值"),
                     prompt_name: str = Body("default",
                                             description="使用的prompt模板名称(在configs/prompt_config.py中配置)"),
                     ):
    current_timestamp = time.time()
    # 实例化Redis连接类
    redis_client = RedisClient(host=REDIS["host"], port=REDIS["port"], db=REDIS["db"]["XiaoKa"],
                               max_connections=REDIS["max_connections"])
    # 实例化对话历史处理类
    chat_history = ChatHistory(uid, redis_client, HISTORY_LEN)
    history = [History.from_data(h) for h in history]

    async def agent_chat_iterator(
            query: str,
            history: Optional[List[History]],
            model_name: str = LLM_MODELS[0],
            prompt_name: str = prompt_name,
    ) -> AsyncIterable[str]:
        nonlocal max_tokens
        callback = CustomAsyncIteratorCallbackHandler()
        if isinstance(max_tokens, int) and max_tokens <= 0:
            max_tokens = None

        model = get_ChatOpenAI(
            model_name=model_name,
            temperature=temperature,
            max_tokens=max_tokens,
            callbacks=[callback],
        )

        kb_list = {x["kb_name"]: x for x in get_kb_details()}
        model_container.DATABASE = {name: details['kb_info'] for name, details in kb_list.items()}

        if Agent_MODEL:
            model_agent = get_ChatOpenAI(
                model_name=Agent_MODEL,
                temperature=temperature,
                max_tokens=max_tokens,
                callbacks=[callback],
            )
            model_container.MODEL = model_agent
        else:
            model_container.MODEL = model

        prompt_template = get_prompt_template("agent_chat", prompt_name)
        prompt_template_agent = CustomPromptTemplate(
            template=prompt_template,
            tools=tools,
            input_variables=["input", "intermediate_steps", "history"]
        )
        output_parser = CustomOutputParser()
        llm_chain = LLMChain(llm=model, prompt=prompt_template_agent)
        memory = ConversationBufferWindowMemory(k=HISTORY_LEN * 2)
        for message in history:
            if message.role == 'user':
                memory.chat_memory.add_user_message(message.content)
            else:
                memory.chat_memory.add_ai_message(message.content)
        if "chatglm3" in model_container.MODEL.model_name or "zhipu-api" in model_container.MODEL.model_name:
            agent_executor = initialize_glm3_agent(
                llm=model,
                tools=tools,
                callback_manager=None,
                prompt=prompt_template,
                input_variables=["input", "intermediate_steps", "history"],
                memory=memory,
                verbose=True,
            )
        else:
            agent = LLMSingleActionAgent(
                llm_chain=llm_chain,
                output_parser=output_parser,
                stop=["\nObservation:", "Observation"],
                allowed_tools=tool_names,
            )
            agent_executor = AgentExecutor.from_agent_and_tools(agent=agent,
                                                                tools=tools,
                                                                verbose=True,
                                                                memory=memory,
                                                                )
        while True:
            try:
                task = asyncio.create_task(wrap_done(
                    agent_executor.acall(query, callbacks=[callback], include_run_info=True),
                    callback.done))
                break
            except:
                pass

        if stream:
            answer_in = True
            async for chunk in callback.aiter():

                # Use server-sent-events to stream the response
                data = json.loads(chunk)
                llm_token = data["llm_token"]
                if data["status"] == Status.start or data["status"] == Status.complete:
                    continue
                elif answer_in:
                    for words in ["Answer:", "Answer"]:
                        index = llm_token.find(words)
                        if index != -1:
                            if words == "Answer:":
                                if llm_token[index + len(words):].lstrip():
                                    yield jsonify(ChatResponseCompletions(code=200, messages=ChatMessageContent(
                                        response=llm_token[index + len(words):], instruction="-1", unique_id=unique_id,
                                        finish=False, data={"content": ""})))
                            answer_in = False

                            break
                elif data["status"] == Status.agent_finish:
                    final_answer = data["final_answer"]
                    chat_history.add_to_history([{"role": "user", "content": query},
                                                 {"role": "assistant", "content": final_answer},
                                                 current_timestamp])
                    log_content(uid=uid, query=query, answer=final_answer, instruction="-1")
                    yield jsonify(ChatResponseCompletions(code=200,
                                                          messages=ChatMessageContent(response=final_answer,
                                                                                      instruction="-1",
                                                                                      unique_id=unique_id,
                                                                                      finish=True, data={"content": ""})))
                else:
                    yield jsonify(ChatResponseCompletions(code=200, messages=ChatMessageContent(response=llm_token,
                                                                                                instruction="-1",
                                                                                                unique_id=unique_id,
                                                                                                finish=False, data={"content": ""})))


        else:
            final_answer = ""
            async for chunk in callback.aiter():
                data = json.loads(chunk)
                if data["status"] == Status.start or data["status"] == Status.complete:
                    continue
                if data["status"] == Status.agent_finish:
                    final_answer = data["final_answer"]

            yield jsonify(ChatResponseCompletions(code=200,
                                                  messages=ChatMessageContent(response=final_answer, instruction="-1",
                                                                              unique_id=unique_id, finish=True, data={"content": ""})))
        await task

    return agent_chat_iterator(query=query,
                               history=history,
                               model_name=model_name,
                               prompt_name=prompt_name)
